import torch
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm
import argparse
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
import subprocess

# python rvm_tool.py --input-video "C:\Users\Missi\Videos\素材\dao.mp4" --output-file "C:\Users\Missi\Videos\素材\dao.gif" --output-format "gif" --model "resnet50"

class VideoReader:
    def __init__(self, path):
        self.cap = cv2.VideoCapture(path)
        self.frame_rate = self.cap.get(cv2.CAP_PROP_FPS)
        if not self.frame_rate or self.frame_rate == 0:
            print("警告: 无法从视频文件中读取有效的帧率，将使用默认值 30 FPS。")
            self.frame_rate = 30
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.num_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.transform = ToTensor()

    def __len__(self):
        return self.num_frames

    def __iter__(self):
        while self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                break
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            yield self.transform(frame_rgb)
        self.cap.release()

class VideoWriter:
    def __init__(self, path, frame_rate, width, height):
        self.writer = cv2.VideoWriter(
            path,
            cv2.VideoWriter_fourcc(*'mp4v'),
            frame_rate,
            (width, height)
        )

    def write(self, frame_tensor):
        frame_np = frame_tensor.mul(255).byte().cpu().numpy().transpose(1, 2, 0)
        frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR)
        self.writer.write(frame_bgr)

    def release(self):
        self.writer.release()


def main(args):
    # 选择设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"将使用设备: {device}")

    # 加载模型
    print(f"正在加载RVM模型 ({args.model})...")
    model = torch.hub.load("PeterL1n/RobustVideoMatting", args.model)
    model = model.to(device)
    model.eval()
    print("模型加载完成！")

    # 准备输入
    reader = VideoReader(args.input_video)

    # 初始化循环状态
    rec = [torch.zeros(1, 1, 1, 1, device=device)] * 4  # for the 4 recurrent states, initialized on the correct device
    downsample_ratio = 0.25 # a hyperparameter, can be adjusted

    if args.output_format == 'gif':
        # --- 通过 FFmpeg 生成透明背景的 GIF ---
        if args.background_image or args.background_color != '0,0,0':
            print("警告：生成GIF时将忽略背景设置，始终输出透明背景。")

        print("准备通过 FFmpeg 生成 GIF...")
        command = [
            'ffmpeg',
            '-y',  # Overwrite output file if it exists
            '-f', 'rawvideo',
            '-vcodec', 'rawvideo',
            '-pix_fmt', 'rgba',
            '-s', f'{reader.width}x{reader.height}',
            '-r', str(reader.frame_rate),
            '-i', '-',  # The input comes from a pipe
            '-vf', 'split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse=dither=bayer:bayer_scale=5:diff_mode=rectangle',
            args.output_file
        ]
        
        process = subprocess.Popen(command, stdin=subprocess.PIPE)
        
        with torch.no_grad():
            for src in tqdm(reader, total=len(reader), desc="生成GIF帧"):
                src = src.to(device).unsqueeze(0)
                # The recurrent states `rec` are already on the correct device
                fgr, pha, *rec = model(src, *rec, downsample_ratio=downsample_ratio)
                
                # 合并前景和透明度通道为 RGBA
                rgba = torch.cat([fgr, pha], dim=1)
                
                # 转换为numpy数组并写入管道
                output_frame_np = rgba.squeeze(0).mul(255).byte().cpu().numpy().transpose(1, 2, 0)
                process.stdin.write(output_frame_np.tobytes())
        
        process.stdin.close()
        return_code = process.wait()
        if return_code == 0:
            print(f"处理完成！GIF 已保存至: {args.output_file}")
        else:
            print(f"\nFFmpeg 处理失败，返回码: {return_code}")
            print("请检查您的FFmpeg是否已正确安装并添加到系统路径中。")

    else: # output_format == 'video'
        # --- 生成视频文件 ---
        writer = VideoWriter(
            args.output_file, 
            reader.frame_rate, 
            reader.width, 
            reader.height
        )

        # 准备背景
        if args.background_image:
            print(f"使用图片 '{args.background_image}' 作为背景。")
            bg_image = Image.open(args.background_image).convert("RGB")
            bg_image = bg_image.resize((reader.width, reader.height))
            background = ToTensor()(bg_image).unsqueeze(0).to(device)
        else:
            print(f"使用纯色 '{args.background_color}' 作为背景。")
            r, g, b = map(int, args.background_color.split(','))
            background = torch.tensor([r/255., g/255., b/255.], device=device).view(1, 3, 1, 1)

        # 处理视频
        print("开始处理视频...")
        with torch.no_grad():
            for src in tqdm(reader, total=len(reader), desc="生成视频帧"):
                src = src.to(device).unsqueeze(0) # add batch dimension
                
                # The recurrent states `rec` are already on the correct device
                fgr, pha, *rec = model(src, *rec, downsample_ratio=downsample_ratio)
                
                # 合成
                com = fgr * pha + background * (1 - pha)
                
                writer.write(com.squeeze(0)) # remove batch dimension

        writer.release()
        print(f"处理完成！视频已保存至: {args.output_file}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="使用Robust Video Matting去除视频背景")
    parser.add_argument("--input-video", required=True, help="输入视频文件路径")
    parser.add_argument("--output-file", required=True, help="输出文件路径")
    parser.add_argument("--model", choices=['mobilenetv3', 'resnet50'], default='mobilenetv3', help="选择要使用的模型。'resnet50'更精确但更慢。")
    parser.add_argument("--output-format", choices=['video', 'gif'], default='video', help="输出格式： 'video' (mp4) 或 'gif' (透明背景动图)")
    
    # 视频专用参数
    parser.add_argument("--background-image", help="用作背景的图片路径（仅用于video格式）")
    parser.add_argument("--background-color", default="0,0,0", help="用作背景的RGB颜色，格式 'R,G,B'（仅用于video格式，默认为黑色）")
    
    args = parser.parse_args()
    main(args) 